#############################################################################
# GFE
#
#############################################################################

rm(list=ls())
graphics.off()

library(data.table)
library(ggplot2)
library(haven)
library(xtable)
library(fixest)
library(mvtnorm)
library(Hmisc)
library("RColorBrewer")

t0 =Sys.time()
memory.limit(size=200000)

#############################################################################
# 1. Data and generate initial values from FE
#############################################################################

# Bring data and FE results
load("Data_new/FE_results.RData")
rm(list=setdiff(ls(),list("dt","dt_res","f3","t0")))
names(dt)

# Generate state-year dummies
state_list = sort(unique(dt$state))
state_year_list = c()
for (s in seq(1,length(state_list),1)) {
  # Per state, generate one dummy per year, except for the last year
  year_now = sort(unique(dt[state==state_list[s],year]))
  for (y in seq(1,length(year_now)-1,1)) {
    dt[,vnow:=as.numeric(state==state_list[s] & year==year_now[y])]
    aux=max(dt$vnow)
    if (aux==0) {
      dt[,vnow:=NULL]
    } else if (aux!=0) {
      setnames(dt,"vnow",paste0("sy_",state_list[s],"_",year_now[y]))
      state_year_list = c(state_year_list,paste0("sy_",state_list[s],"_",year_now[y]))
    }
  }
}
summary(dt$sy_1_1950)
length(state_year_list)

# Re-estimate the model again to get the coefficients of the state-year FE
form=as.formula(paste("yield ~ precr + ddr:factor(fips) +", 
                      paste(state_year_list, collapse= "+")," | factor(fips)"))
form
fe = feols(fml=form,data=dt,cluster="state")

# Save common weather and state-year FE coefficients
gamma0=fe$coefficients[c("precr",state_year_list)]
summary(gamma0)

# Save slopes and compare with the originally estimated ones (rounding errors e-10)
source("Code/0. CODE Auxiliary.R")
dt_slop = slopes_dt(f=fe,names_var="")
dim(dt_slop)
summary(dt_slop)
dt_slop = merge(dt_slop,dt_res[,lapply(.SD,min),.SDcols="mg3",by="fips"],by="fips",all=TRUE)
dim(dt_slop)
summary(dt_slop)
summary(dt_slop[,coef-mg3])
dt_slop[,mg3:=NULL]

# Save county FE 
dt_fe=data.table(fips=as.numeric(names(fixef(fe)$`factor(fips)`)),alpha=fixef(fe)$`factor(fips)`)
dim(dt_fe)
summary(dt_fe)

### Use the SE as a baseline var for the looping over starting values
# For beta, mean Std. Error is smaller than sd(beta_i). Keep sd(beta_i)
summary(dt_slop)
sd(dt_slop$coef)
mean(dt_slop$se)
se_beta = sd(dt_slop$coef)

# For gamma, keep vector of Std. Errors x 10
summary(fe$coeftable)
fe$coeftable[1:2,]
se_gamma = fe$coeftable[names(gamma0),"Std. Error"]*10 

# Demeaned data at the county level
cols_old=c("yield","ddr","precr",state_year_list)
dt_dm = dt[,.SD-lapply(.SD,mean),.SDcols=cols_old,by="fips"]  
summary(dt_dm)
dim(dt_dm)

# Add a correlative obs number within county
dt_dm[,ones:=1]
dt_dm[,nf:=cumsum(ones),by="fips"]
dt_dm[,ones:=NULL]

setnames(dt_slop,"coef","slope")
rm(aux,year_now,y,s,state_list)
rm(fe,form,cols_old)

t1 =Sys.time()
t1-t0
rm(t0,t1)
save.image("Data_new/GFE/Data_for_groups.RData")


#############################################################################
# 2. Run several grouped FE
#############################################################################

load("Data_new/GFE/Data_for_groups.RData")

# Source the right functions
source("Code/5a. CODE grouped FE.R")

niter=10000
tol=1e-7
sseed=20230711
save.image("Data_new/GFE/Data_for_groups.RData")

######################################
# G=2
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=2
Ns=3000
save_name="g2"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
                se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
                sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

######################################
# G=3
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=3
Ns=3000
save_name="g3"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
                se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
                sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

######################################
# G=4
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=4
Ns=3000
save_name="g4"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
            se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
            sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

######################################
# G=5
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=5
Ns=3000
save_name="g5"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
            se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
            sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

######################################
# G=6
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=6
Ns=3000
save_name="g6"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
            se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
            sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

######################################
# G=7
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=7
Ns=3000
save_name="g7"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
                se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
                sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

######################################
# G=8
######################################

load("Data_new/GFE/Data_for_groups.RData")
G=8
Ns=3000
save_name="g8"

res=gfe_m1_wrap(data=dt_dm,data_ori=dt,G=G,gamma0=gamma0,beta0=dt_slop,
                se_gamma=se_gamma,se_beta=se_beta,Ns=Ns,niter=niter,tol=tol,
                sseed=sseed,save_name=save_name)
save.image(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

#############################################################################
# 3. Review the results one by one, put together
#############################################################################

rm(list=ls())
names_keep = c()

######################################
# G=2
######################################

save_name="g2"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g2 = res$res
dc_g2 = res$dc
time_g2 = res$time
names_keep= c(names_keep,save_name)

######################################
# G=3
######################################

save_name="g3"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g3 = res$res
dc_g3 = res$dc
time_g3 = res$time
names_keep= c(names_keep,save_name)

######################################
# G=4
######################################

save_name="g4"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g4 = res$res
dc_g4 = res$dc
time_g4 = res$time
names_keep= c(names_keep,save_name)

######################################
# G=5
######################################

save_name="g5"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g5 = res$res
dc_g5 = res$dc
time_g5 = res$time
names_keep= c(names_keep,save_name)

######################################
# G=6
######################################

save_name="g6"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g6 = res$res
dc_g6 = res$dc
time_g6 = res$time
names_keep= c(names_keep,save_name)

######################################
# G=7
######################################

save_name="g7"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g7 = res$res
dc_g7 = res$dc
time_g7 = res$time
names_keep= c(names_keep,save_name)

######################################
# G=8
######################################

save_name="g8"
load(paste0("Data_new/GFE/Gfe_res1_",save_name,".RData"))

# Save with a new name to not overwrite
res_g8 = res$res
dc_g8 = res$dc
time_g8 = res$time
names_keep= c(names_keep,save_name)

######################################
# Keep only results and save
######################################

names_keep
names_keept = c(paste0("res_",names_keep),paste0("dc_",names_keep),paste0("time_",names_keep))
names_keept = c(names_keept,"dt_dm","dt","gamma0","dt_slop","se_gamma","se_beta","niter","tol","sseed","names_keep")

rm(list=setdiff(ls(),names_keept))

# Source the right functions
source("Code/5a. CODE grouped FE.R")

save_res = "Results"
save.image("Data_new/GFE/GFE_results.RData")

